"""
Copyright 2020 Huawei Technologies Co., Ltd

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import os
import json
import argparse
import traceback
from common import set_env
from common import atc_model
from flat_generate_model import get_dist_compute_json

# number of distance accumulate per time
dist_accum_num = 2048
nprobe_slice = 8


def arg_parse():
    """
    Parse arguements to the operator model
    """

    parser = argparse.ArgumentParser(
        description='generate aicore operator model')

    parser.add_argument("-n",
                        dest='queries_num',
                        default=1,
                        type=int,
                        help="Number of query vector")
    parser.add_argument("-d",
                        dest='dim',
                        default=512,
                        type=int,
                        help="Feature dimension")
    parser.add_argument("-c",
                        dest='coarse_centroid_num',
                        default=16384,
                        type=int,
                        help="Number of coarse centroid")
    parser.add_argument("-s",
                        dest='sub_quantizers_num',
                        default=64,
                        type=int,
                        help="Number of sub-quantizers per vector")
    parser.add_argument("-q",
                        dest='pq_centroids_num',
                        default=256,
                        type=int,
                        help="Number of pq centroids")

    return parser.parse_args()


def get_dist_table_build_json(dim, coarse_centroid_num, sub_quantizers_num,
                              pq_centroids_num, file_path):
    # write dist_table_build json
    dist_table_build_obj = [{
        "op":
        "DistanceTableBuild",
        "input_desc": [{
            "format": "ND",
            "shape": [1, dim],
            "type": "float16"
        }, {
            "format":
            "ND",
            "shape":
            [sub_quantizers_num, pq_centroids_num * dim // sub_quantizers_num],
            "type":
            "float16"
        }, {
            "format": "ND",
            "shape": [1, nprobe_slice],
            "type": "int32"
        }, {
            "format": "ND",
            "shape": [coarse_centroid_num, dim],
            "type": "float16"
        }],
        "output_desc": [{
            "format":
            "ND",
            "shape": [1, nprobe_slice, sub_quantizers_num, pq_centroids_num],
            "type":
            "float16"
        }, {
            "format": "ND",
            "shape": [16],
            "type": "uint16"
        }]
    }]

    isObj = json.dumps(dist_table_build_obj)
    fileObject = open(file_path, 'w')
    fileObject.write(isObj)
    fileObject.close()


def get_dist_accum_json(sub_quantizers_num, file_path):
    # write dist_accum json
    dist_accum_obj = [{
        "op":
        "DistAccum",
        "input_desc": [{
            "format":
            "ND",
            "shape": [dist_accum_num, (sub_quantizers_num + 15) // 16 * 16],
            "type":
            "float16"
        }],
        "output_desc": [{
            "format": "ND",
            "shape": [dist_accum_num],
            "type": "float16"
        }, {
            "format": "ND",
            "shape": [16],
            "type": "uint16"
        }]
    }]

    isObj = json.dumps(dist_accum_obj)
    fileObject = open(file_path, 'w')
    fileObject.write(isObj)
    fileObject.close()


def generate_by_atc():
    # generate directory of model
    dirs = './op_models/'
    if not os.path.exists(dirs):
        os.makedirs(dirs)

    atc_model('distance_compute_op1')
    atc_model('distance_compute_op2')
    atc_model('distance_compute_op4')
    atc_model('distance_compute_op8')
    atc_model('distance_compute_op16')
    atc_model('distance_compute_op32')
    atc_model('dist_table_build_op')
    atc_model('dist_accum_op')


if __name__ == '__main__':
    args = arg_parse()
    set_env()

    queries_num = args.queries_num
    dim = args.dim
    coarse_centroid_num = args.coarse_centroid_num
    sub_quantizers_num = args.sub_quantizers_num
    pq_centroids_num = args.pq_centroids_num

    coarse_centroid_num = (coarse_centroid_num + 15) // 16 * 16
    nprobe_slice = (nprobe_slice + 7) // 8 * 8

    # generate directory of config
    dirs = './config/'
    if not os.path.exists(dirs):
        os.makedirs(dirs)

    get_dist_compute_json(1, dim, coarse_centroid_num,
                          './config/distance_compute_op1.json')
    get_dist_compute_json(2, dim, coarse_centroid_num,
                          './config/distance_compute_op2.json')
    get_dist_compute_json(4, dim, coarse_centroid_num,
                          './config/distance_compute_op4.json')
    get_dist_compute_json(8, dim, coarse_centroid_num,
                          './config/distance_compute_op8.json')
    get_dist_compute_json(16, dim, coarse_centroid_num,
                          './config/distance_compute_op16.json')
    get_dist_compute_json(32, dim, coarse_centroid_num,
                          './config/distance_compute_op32.json')

    get_dist_table_build_json(dim, coarse_centroid_num, sub_quantizers_num,
                              pq_centroids_num,
                              './config/dist_table_build_op.json')

    get_dist_accum_json(sub_quantizers_num, './config/dist_accum_op.json')

    try:
        generate_by_atc()
    except OSError:
        traceback.print_exc()
        print("make sure set right ASCEND_HOME environment variable")
    else:
        traceback.print_exc()
